execfile("addcon_mapping.py")

from pysrt import SubRipFile, SubRipItem, SubRipTime
import os, sys, random, re, shutil
from addcon_script import Script, LineWord

path_to_srt = os.path.abspath(sys.argv[1])
path_to_script = os.path.abspath(sys.argv[2])
save_to_path = path_to_srt
shutil.copyfile(path_to_srt, path_to_srt+".back")
subs = SubRipFile.open(path_to_srt)


ACTOREXP = r'^(?P<ACTOR>\w+:\s)?'
LINESEXP = r'(?P<LINE>[^\[^\]]+)'
CONTEXTEXP = r'(?P<CONTEXT>\[[^\[^\]]+\])'
SUBS_PRESERVE = 0 # if you preserve original subtitiles 

COMBINEEXP = re.compile( LINESEXP + "|" + CONTEXTEXP )

scripts = []

masterindex = 1
for iScript, line in enumerate(open(path_to_script, 'r')):
    ma = re.match(ACTOREXP + "(.+)", line.strip())
    if ma:
        if ma.group(2):
            for linecntxt in re.findall(COMBINEEXP,ma.group(2)):
                if linecntxt[0] == "":
                    "Context"
                    scripts.append(Script(linecntxt[1],iScript,"Context", masterindex, subindex = []))
                    masterindex += 1
                elif linecntxt[1] == "":
                    "Line"
                    scripts.append(Script(linecntxt[0],iScript,"Line", masterindex, subindex = []))
                    masterindex += 1

words_script = []
words_subtitle = []
for sc in scripts:
    if sc.attr == "Line":
        #print "This is Line" , sc.masterindex, sc.script_index
        # Break Lines
        # remove "!?."
        for word in re.sub("[\(\)\\?\!\.\"\,]+", "", sc.text).strip().split():
            words_script.append(LineWord(word, sc.masterindex, None))


for sub in subs:
    for word in re.sub("[\(\)\\?\!\.\"\,]+", "", sub.text).strip().split():
        words_subtitle.append(LineWord(word, None, sub.index))



def levenshtein_mapping(a,b):
    """
    original: http://d.hatena.ne.jp/naoya/20090329/1238307757
    author: funaya@gmail.com, hiroyuki-fn in Twitter
    
    input: two lists: lista and list b
    output: mapping of which elements in a corresponds to which elements in b
            in terms of Levenshtein distance
    
    """
    m = [ [0] * (len(b) + 1) for i in range(len(a) + 1) ]

    for i in xrange(len(a) + 1):
        m[i][0] = i

    for j in xrange(len(b) + 1):
        m[0][j] = j

    for i in xrange(1, len(a) + 1):
        if i%100 == 0 :print i, " / ", len(a)
        word1 = a[i - 1].word
        for j in xrange(1, len(b) + 1):
            if word1 == b[j - 1].word:
                x = 0
            else:
                x = 1
            m[i][j] = min(m[i - 1][j] + 1, m[i][ j - 1] + 1, m[i - 1][j - 1] + x)
            
    # calculating mapping
    lmap = []
    i = len(a)
    j = len(b)
    
    while i >= 1 and j >= 1:
        candidates = [m[i - 1][j - 1], m[i - 1][j] , m[i][j - 1] ]
        next_distance = min(candidates)
        cani = [[i - 1, j - 1], [i - 1 , j], [i , j - 1]]
        
        if ( m[i][j] == next_distance ) and ( candidates.index(next_distance) == 0) :
            lmap.append([i-1, j-1])
            
        i, j = cani[candidates.index(next_distance)]
           
    lmap.reverse() 
    return lmap

print "Calculating Levenshtein..."
map = levenshtein_mapping(words_script, words_subtitle)


# Natural Insertion (linearly inserted when data is missing)
mprev = None
mapidx = 1
mprev = map[0]
while mapidx < len(map):
    m = map[mapidx]
    dif1 = m[0] - mprev[0]
    dif2 = m[1] - mprev[1]
    if dif2 >= 2 and dif1 >= 2:
        ni = max(dif1,dif2) -1 # number of insertions
        for j in range(0, ni):
            map.insert(mapidx + j, [map[mapidx-1][0] + (j + 1)*(dif1-1)/ni, map[mapidx-1][1] + (j + 1)*(dif2-1)/ni, "added"])
            print "inserted at %d" % mapidx
        mapidx += dif1 -1
    mapidx +=1
    mprev = m
# print map




for m in map:
    words_script[m[0]].subindex = words_subtitle[m[1]].subindex
    # words_subtitle[m[1]].masterindex = words_subtitle[m[0]].scriptindex
    

print "adding subindex..."
for sc in scripts:
    # adding possible subindex to all lines in scripts
    if sc.attr == "Line":
        list = [ws for ws in words_script if ws.masterindex == sc.masterindex]
        for ws in list:
            if ws.subindex != None:
                sc.subindex.append(ws.subindex)

number_context = 0
number_assigned = 0
postlist =[]
assignedlist =[]
for i,sc in enumerate(scripts):

    prev, next = [None, None]
    prevs, nexts = [None, None] # valid if it is in the same line in scripts

    if sc.attr == "Line":
        # print sc.subindex
        continue
    else: # sc.attr == "Context":
        number_context +=1
        
        if i != len(scripts)-1:
            if scripts[i+1].subindex: # scripts[1].attr == "Line":
                next = scripts[i+1].subindex[0]
                if sc.script_index == scripts[i+1].script_index:
                    nexts = next
        if i != 0:
            if scripts[i-1].subindex:
                prev = scripts[i-1].subindex[-1]
                if sc.script_index == scripts[i-1].script_index:
                    prevs = prev
                    
        # print prevs, prev, next, nexts, sc.text
        
        # inserting context to subs
        if prevs == None and nexts == None:
            # It's probably a description of a scene
            if prev and next == None:
                start_min = subs[prev-1].end
                end_max = subs[prev].start
                subs[prev-1].start = start_min
                subs[prev-1].end = end_max
                subs[prev-1].text = subs[prev-1].text * SUBS_PRESERVE + sc.text + "\n"
                number_assigned += 1
                assignedlist.append(i)
                continue
            elif next:
                if next-2 >= 0:
                    start_min = subs[next-2].end # so that no crossover
                end_max = subs[next].start # so that no crossover
                ratio = float(len(sc.text+subs[next-1].text))/len(subs[next-1].text)
                duration = (subs[next-1].end - subs[next-1].start)*ratio
                start_tmp = subs[next-1].end - duration
                
                subs[next-1].text = sc.text + "\n" + subs[next-1].text * SUBS_PRESERVE
                if start_tmp < start_min:
                    subs[next-1].start = start_min
                else:
                    subs[next-1].start = start_tmp
                number_assigned += 1
                assignedlist.append(i)
                
            elif prev == None and next == None:
                # Insert New Subtitles Later
                postlist.append(i)
                continue
                
        elif prevs and nexts == None:
            subs[prevs-1].text = subs[prevs-1].text * SUBS_PRESERVE + sc.text + "\n" 
            number_assigned += 1
            assignedlist.append(i)
        elif nexts:
            subs[nexts-1].text = sc.text + "\n"  + subs[nexts-1].text * SUBS_PRESERVE
            number_assigned += 1
            assignedlist.append(i)

print "number ", number_assigned, " of ", number_context, " assigned." 
print "%d" % len(postlist), " unussigned"
print  assignedlist

subs.save( save_to_path ,"utf-8")
print "saved to %s." % save_to_path
